import json
import os
import time
import tracemalloc

from flask import Flask, request

from FaceRec.DBFace.UseDBFace import test_detector, test_landmarks, reload_weights
from FaceRec.DBFace.utils import b64string2array

tracemalloc.start()
app = Flask(__name__)



@app.route('/imr-ai-service/atomic_functions/faces_detect', methods=['POST'])
def faces_detect():
    snapshot = tracemalloc.take_snapshot()  # 快照，当前内存分配
    top_stats = snapshot.statistics('lineno')  # 快照对象的统计
    for stat in top_stats:
        print(stat)
    if request.method == "POST":
        c_da = request.data
        data = json.loads(c_da.decode())
        img_string = data['imgString']
        img = b64string2array(img_string)
        time_take = time.time()
        if img is None:
            result = []
        else:
            result = test_detector(img)
        time_take = time.time() - time_take
        # os.remove(path)
        snapshot = tracemalloc.take_snapshot()  # 快照，当前内存分配
        top_stats = snapshot.statistics('lineno')  # 快照对象的统计
        for stat in top_stats:
            print(stat)
        return json.dumps({'res': result, 'timeTake': round(time_take, 4)},
                          ensure_ascii=False)


@app.route('/imr-ai-service/atomic_functions/landmarks_detect', methods=['POST'])
def landmarks_detect():
    if request.method == "POST":
        c_da = request.data
        data = json.loads(c_da.decode())
        img_string = data['imgString']
        img = b64string2array(img_string)
        time_take = time.time()
        if img is None:
            result = []
        else:
            result = test_landmarks(img)
        time_take = time.time() - time_take
        return json.dumps({'res': result, 'timeTake': round(time_take, 4)},
                          ensure_ascii=False)


@app.route('/imr-ai-service/atomic_functions/reload', methods=['POST'])
def reload():
    if request.method == "POST":
        time_take = time.time()
        result = reload_weights()
        time_take = time.time() - time_take
        return json.dumps(
            {
                'total': result[0],
                'new': result[1],
                'timeTake': round(time_take, 4)
            },
            ensure_ascii=False
        )


if __name__ == '__main__':
    pid = os.getpid()
    print('pid is:', pid)
    with open('atomic_pid.txt', 'w') as f:
        f.writelines([str(pid)])
    app.run(
        host="0.0.0.0",
        port=int("12242"),
        debug=False, threaded=True)
